package com.ld.poetry.filter;


import com.ld.poetry.redis.RedisStringService;
import com.ld.poetry.utils.HashUtils;
import com.ld.poetry.utils.RequestUtils;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import javax.servlet.*;
import javax.servlet.annotation.WebFilter;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.util.concurrent.TimeUnit;


/**
 * created by vital on 2024/10/7 下午2:05
 * 新增功能，ip限流过滤器
 * 用于限制ip访问次数
 */
@Slf4j
@Component
@WebFilter(urlPatterns = "/*", filterName = "IpRateLimitFilter")
public class IpRateLimitFilter implements Filter {

    private static final int MAX_REQUESTS_PER_MINUTE = 10; // 相同路径每分钟最大请求次数
    private static final long EXPIRATION = 1; // 过期时间为1
    private static final TimeUnit TIME_UNIT = TimeUnit.MINUTES;// 时间单位：分钟

    @Resource
    private RedisStringService redisStringService;


    @Override
    public void doFilter(ServletRequest servletRequest, ServletResponse servletResponse, FilterChain filterChain) throws IOException, ServletException {
        log.info("IpRateLimitFilter start");

        HttpServletRequest request = (HttpServletRequest) servletRequest;
        HttpServletResponse response = (HttpServletResponse) servletResponse;

        // 直接放行预检请求
        if(request.getMethod().equals("OPTIONS")){
            filterChain.doFilter(servletRequest, servletResponse);
            return;
        }

        // 生成key
        String key = this.generateKey(request);

        // 判断请求体是否是第一次请求
        Boolean hasKey = redisStringService.hasKey(key);
        // 判断是否是第一次访问
        if(!hasKey){
            //第一次访问，初始化请求次数
            redisStringService.setStringWithExpire(key,1,EXPIRATION,TIME_UNIT);
            filterChain.doFilter(servletRequest, servletResponse);
            return;
        }

        //已访问过，判断是否超过最大请求次数
        Boolean result = this.checkRequestLimit(key, response);
        if(!result){
            // 超过最大速率，拒绝访问
            return;
        }

        filterChain.doFilter(servletRequest, servletResponse);
    }

    @Override
    public void init(FilterConfig filterConfig) throws ServletException {
        log.info("IpRateLimitFilter init");
        Filter.super.init(filterConfig);
    }

    @Override
    public void destroy() {
        log.info("IpRateLimitFilter destroy");
        Filter.super.destroy();
    }


    /**
     * 根据请求ip，url,useragent生成key
     * @param request 请求对象
     * @return 生成的key
     */
    private String generateKey(HttpServletRequest request){
        //获取请求ip
        String ip = RequestUtils.getRealIp(request);
        // 获取请求路径
        String uri = request.getRequestURI();
        // 获取请求头中的useragent
        String userAgent = request.getHeader("User-Agent");
        // 合成key
        String key = "ip_" + ip + "_" + uri + "_" + userAgent ;
        // 压缩key
        return HashUtils.sha256(key);
    }

    /**
     * 检查请求次数是否超过限制
     * @param key generateKey方法合成的key
     * @param response 响应对象
     * @return 超过限制返回false, 否则返回true
     */
    private Boolean checkRequestLimit(String key, HttpServletResponse response) throws IOException {
        // 获取访问次数
        Integer count = (Integer) redisStringService.getString(key);
        //判断是否超过最大请求次数
        if(count >= MAX_REQUESTS_PER_MINUTE){
            //拒绝访问
            response.setStatus(429);
            // 设置内容类型为HTML
            response.setContentType("text/html;charset=UTF-8");

            // 写入响应体 - 简单的HTML页面
            String errorMessage = "<!DOCTYPE html>" +
                    "<html>" +
                    "<head><title>请求频繁</title></head>" +
                    "<body>" +
                    "<h1>请求频繁</h1>" +
                    "<p>您在短时间内发起了太多的请求，请5分钟后再试。</p>" +
                    "</body>" +
                    "</html>";
            response.getWriter().write(errorMessage);

            // 封禁5分钟
            redisStringService.setExpires(key,5,TIME_UNIT);
            log.warn("key:{}。{}分钟内，访问次数超过" + MAX_REQUESTS_PER_MINUTE + "次，已封禁5分钟", key, EXPIRATION);

            // 更新ip请求次数
            redisStringService.increment(key);
            return Boolean.FALSE;
        }
        // 更新ip请求次数
        redisStringService.increment(key);
        return Boolean.TRUE;
    }

}
